#DESCRIPTION:
#generic functions for plotting data in stereonet


#1. libraries and modules

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.stats import beta
import matplotlib.cm as cm




def spherical_to_polar_to_cart(dip,ddir):
    #spherical (dip-dip direction or plunge-trend)
    dip=np.asarray(dip)
    ddir=np.asarray(ddir)
    z=1+np.sin(np.radians(dip))
    x=np.cos(np.radians(dip))
    zenith_angle=(np.arctan(x/z))
    #to polar  (radius, theta)
    r=np.tan(zenith_angle)
    theta=-ddir+90
    #to cartesian (x,y)
    X=r*np.cos(np.radians(theta))
    Y=r*np.sin(np.radians(theta))
    return X,Y

def plot_stereonet():
    #plot blank stereonet with small circles and E-W and N-S lines
    small_circles=[15,30,45,60,75]
    theta=np.linspace(0,2*np.radians(360),1000)
    ax=plt.gca()
    X=np.cos(theta)
    Y=np.sin(theta)
    ax.plot(X,Y,'k-')
    ax.plot([-1,1],[0,0],'k-',lw=0.1)
    ax.plot([0,0],[-1,1],'k-',lw=0.1)
    sc_list=[]
    for d in small_circles:
        z=1+np.sin(np.radians(d))
        x=np.cos(np.radians(d))
        zenith_angle=(np.arctan(x/z))
        r=np.tan(zenith_angle)
        X=r*np.cos(theta)
        Y=r*np.sin(theta)
        sc=plt.plot(X,Y,'k-',lw=0.1)
        sc_list.append(sc)
    ax.axis('equal')
    ax.set_xticks([])
    ax.set_yticks([])
    return ax,sc_list

def draw_great_circle(dip,ddir,lc,lw):
    #plot great circle on stereonet
    a=np.linspace(-90,90,1000)
    adipdir_list=ddir+a

    delta_az=np.abs(adipdir_list-ddir)
    delta_az=np.where(delta_az>180,360-delta_az,delta_az)
    d_adip=np.where(delta_az<=90,
                    np.degrees(np.arctan(np.tan(np.radians(dip))*np.cos(np.radians(delta_az)))),
                    -1*np.ones(delta_az.shape))

    X,Y=spherical_to_polar_to_cart(d_adip,adipdir_list)
    ax=plt.gca()
    gc=ax.plot(X,Y,'-',color=lc,lw=lw)
    return gc

def plot_colorvalue_data(dip,ddir,col,mincol,maxcol,size):
    #plot data with color value in current stereoplot
    X,Y=spherical_to_polar_to_cart(dip,ddir)
    sc=plt.scatter(X[np.argsort(col)],Y[np.argsort(col)],
                  c=np.sort(col),s=size,vmin=mincol,vmax=maxcol,cmap=cm.jet,linewidth=0)
    
    

def plot_data(dip,ddir,size,mar):
    #plot data in current stereoplot
    X,Y=spherical_to_polar_to_cart(dip,ddir)
    ax=plt.gca()
    ax.scatter(X,Y,s=size,marker=mar)
        
    
    
    
